package com.utooo.bid.config;

import com.utooo.bid.service.Assistant;
import com.utooo.bid.service.DocumentIngestionService;
import com.utooo.bid.service.RagAssistant;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.splitter.DocumentSplitters;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.memory.chat.ChatMemoryProvider;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.chat.StreamingChatLanguageModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.pgvector.PgVectorEmbeddingStore;
import lombok.RequiredArgsConstructor;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Scope;

import static org.springframework.beans.factory.config.BeanDefinition.SCOPE_PROTOTYPE;

@Configuration
@RequiredArgsConstructor
public class LangChain4jConfig {


    final PgConfig pgConfig;
    final RedisChatMemoryStore redisChatMemoryStore;

    @Bean
    public EmbeddingStore<TextSegment> embeddingStore() {
        return PgVectorEmbeddingStore.builder()
                .table(pgConfig.getTable())
                .createTable(true)
                .host(pgConfig.getHost())
                .port(pgConfig.getPort())
                .user(pgConfig.getUser())
                .dimension(1024)
                .password(pgConfig.getPassword())
                .database(pgConfig.getDatabase())
                .build();
    }

    @Bean
    @ConditionalOnMissingBean
    public DocumentSplitter documentSplitter() {
        return DocumentSplitters.recursive(30, 20);
    }

    @Bean
    @ConditionalOnMissingBean
    public ContentRetriever contentRetriever(EmbeddingStore<TextSegment> embeddingStore,
                                             EmbeddingModel embeddingModel) {
        return EmbeddingStoreContentRetriever.builder()
                .embeddingStore(embeddingStore)
                .embeddingModel(embeddingModel)
                .maxResults(3)
                .minScore(0.6)
                .build();
    }

    @Bean
    public DocumentIngestionService documentIngestionService(
            DocumentSplitter splitter,
            EmbeddingModel embeddingModel,
            EmbeddingStore<TextSegment> embeddingStore) {
        return new DocumentIngestionService(splitter, embeddingModel, embeddingStore);
    }

    @Bean
    public RagAssistant ragAssistant(ChatLanguageModel chatModel,
                                     StreamingChatLanguageModel streamingChatModel,
                                     ContentRetriever contentRetriever,
                                     ChatMemoryProvider chatMemoryProvider)  {
        return AiServices.builder(RagAssistant.class)
                .chatLanguageModel(chatModel)
                .streamingChatLanguageModel(streamingChatModel)
                .contentRetriever(contentRetriever)
                .chatMemoryProvider(chatMemoryProvider)
                .chatMemory(MessageWindowChatMemory.withMaxMessages(10))
                .build();
    }

    @Bean
    public ChatMemoryProvider chatMemoryProvider() {
        return memoryId -> MessageWindowChatMemory.builder()
                // 设置消息窗口 ID
                .id(memoryId)
                // 设置消息最大条数，默认为 10
                .maxMessages(10)
                .chatMemoryStore(redisChatMemoryStore)
                .build();
    }

    @Bean
    public Assistant assistant(ChatLanguageModel chatModel,
                               ChatMemoryProvider chatMemoryProvider) {
        return AiServices.builder(Assistant.class)
                .chatLanguageModel(chatModel)
                .chatMemoryProvider(chatMemoryProvider)
                .build();
    }
}
